import torch
from torch import nn


class ImageTransform(nn.Module):
    def __init__(self, seq_len=768, vocab_size=196, embedding_dims=1024):
        super().__init__()
        self.seq_len = seq_len
        self.vocab_size = vocab_size

        # 分 14x14 path（768为一堆） => (1,3,224,224) ==> (1,768,14,14)
        self.patch_layer = nn.Conv2d(3, seq_len, 16, 16)

        # 1、形态：需要在768的形态上+1,所以中间的1为为了768设置的。
        # 2、nn.Parameter 证明CLS会被训练（按照以前设计原理cls_token应该是固定值）
        #    Vit的论文中说，CLS被训练时，实际的训练效率会明显高于固定值的状态。
        self.cls_token = nn.Parameter(torch.randn(1, 1, vocab_size))

        self.embedding = nn.Linear(vocab_size, embedding_dims)

        # 位置编码 (可以在对应位置自适应调整)
        self.position_encode = nn.Parameter(torch.randn(1, seq_len + 1, embedding_dims))

    def forward(self, x):
        x = self.patch_layer(x)  # (1,768,14,14)
        x = x.reshape(-1, self.seq_len, self.vocab_size)  # (1,768,196)

        # 对齐x的批次大小
        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # (1,1,196)
        x = torch.cat([cls_token, x], dim=1)  # (1,769,196)
        x = self.embedding(x)
        return x + self.position_encode